/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * Jan 22, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef __QUERY_ALL_PKGS_FOR_AGT_HPP
#define __QUERY_ALL_PKGS_FOR_AGT_HPP

#include "chainOfResponsability.hpp"
#include "loggerHelper.h"
#include "scanContext.hpp"
#include "socketDBWrapper.hpp"
#include "vulnerabilityScanner.hpp"
#include "wazuhDBQueryBuilder.hpp"

/**
 * @brief Orchestrates queries to the global Wazuh-DB system and initiates package scanning.
 *
 * The `TQueryAllPackages` class is responsible for coordinating the execution of queries within the
 * broader Wazuh ecosystem.
 *
 * @tparam TScanContext scan context type.
 */
template<typename TScanContext = ScanContext,
         typename TAbstractHandler = AbstractHandler<std::shared_ptr<TScanContext>>>
class TQueryAllPackages final : public AbstractHandler<std::shared_ptr<TScanContext>>
{
private:
    std::shared_ptr<TAbstractHandler> m_suborchestration;

public:
    /**
     * @brief Construct a new global fetch object
     *
     * @param packageInsertOrchestration package orchestration instance.
     */
    explicit TQueryAllPackages(std::shared_ptr<TAbstractHandler> packageInsertOrchestration)
        : m_suborchestration(std::move(packageInsertOrchestration))
    {
    }

    /**
     * @brief Handles request and passes control to the next step of the chain.
     *
     * This method processes a request based on the operation type specified in the input data.
     * If the operation type is SingleDeleteAndInsert, it queries a single entry from the Wazuh-DB and processes the
     * response. If the operation type is InsertAll, it executes queries for all packages from the Wazuh-DB and
     * processes the responses. If the operation type is not recognized, it throws a std::runtime_error.
     *
     * @param data A shared pointer to the input data.
     * @return A shared pointer to the result of the request processing.
     *
     * @note This function interacts with a Wazuh-DB socket, processes JSON responses, and creates and handles data
     * using flatbuffers. It may also throw exceptions if certain conditions are not met.
     */
    std::shared_ptr<TScanContext> handleRequest(std::shared_ptr<TScanContext> data) override
    {
        // Instance the socketWrapper for Wazuh-DB
        static SocketDBWrapper wdbSocketWrapper(WDB_SOCKET);

        nlohmann::json response;
        if (data->operationType() == OperationType::SingleDeleteAndInsert)
        {
            // Query for the agent packages
            auto idAgent = std::string {data->agentId()};
            logDebug2(WM_VULNSCAN_LOGTAG, "Re-scanning packages for agent '%s'", idAgent.c_str());
            wdbSocketWrapper.query(WazuhDBQueryBuilder::builder().agentGetPackagesCommand(idAgent).build(), response);
            // Validate the response
            if (!response.empty())
            {
                // The global.db stores the agent version with the prefix "Wazuh " when the keepalive is parsed
                std::string normalizedAgentVersion = Utils::leftTrim(data->agentVersion().data(), "Wazuh ");

                for (const auto& package : response)
                {
                    flatbuffers::FlatBufferBuilder fbBuilder;
                    auto agentInfo =
                        SyscollectorSynchronization::CreateAgentInfoDirect(fbBuilder,
                                                                           idAgent.c_str(),
                                                                           std::string(data->agentIp()).c_str(),
                                                                           std::string(data->agentName()).c_str(),
                                                                           std::string(normalizedAgentVersion).c_str(),
                                                                           std::string(data->agentNodeName()).c_str());
                    auto syscollectorPackage = SyscollectorSynchronization::Createsyscollector_packagesDirect(
                        fbBuilder,
                        package.value("architecture", "").c_str(),
                        package.value("checksum", "").c_str(),
                        package.value("description", "").c_str(),
                        package.value("format", "").c_str(),
                        package.value("groups", "").c_str(),
                        package.value("install_time", "").c_str(),
                        package.value("item_id", "").c_str(),
                        package.value("location", "").c_str(),
                        package.value("multiarch", "").c_str(),
                        package.value("name", "").c_str(),
                        package.value("priority", "").c_str(),
                        package.value("scan_time", "").c_str(),
                        package.value("size", 0),
                        package.value("source", "").c_str(),
                        package.value("vendor", "").c_str(),
                        package.value("version", "").c_str());

                    auto stateMsg = SyscollectorSynchronization::Createstate(
                        fbBuilder,
                        SyscollectorSynchronization::AttributesUnion::AttributesUnion_syscollector_packages,
                        syscollectorPackage.Union());

                    auto syncMsg = SyscollectorSynchronization::CreateSyncMsg(
                        fbBuilder,
                        agentInfo,
                        SyscollectorSynchronization::DataUnion::DataUnion_state,
                        stateMsg.Union());

                    fbBuilder.Finish(syncMsg);

                    std::variant<const SyscollectorDeltas::Delta*,
                                 const SyscollectorSynchronization::SyncMsg*,
                                 const nlohmann::json*>
                        variantData = SyscollectorSynchronization::GetSyncMsg(fbBuilder.GetBufferPointer());

                    // This could throw a exception if there is not information about the os, because the user disable
                    // it from syscollector.
                    try
                    {
                        auto context = std::make_shared<TScanContext>(variantData);
                        m_suborchestration->handleRequest(std::move(context));
                    }
                    catch (const std::exception& e)
                    {
                        logDebug1(WM_VULNSCAN_LOGTAG, "Unable to handle request: %s", e.what());
                    }
                }
            }
            else
            {
                logWarn(WM_VULNSCAN_LOGTAG,
                        "Empty response for agent '%s' in Wazuh-DB 'sys_programs' query",
                        std::string(data->agentId()).c_str());
            }
        }
        else if (data->operationType() == OperationType::InsertAll)
        {
            logDebug2(WM_VULNSCAN_LOGTAG, "Re-scanning packages for all agents");
            for (const auto& agent : data->getMessages())
            {
                // Execute query for all packages, modify to adapt to testtool input.
                std::string agentIdStr;
                if (agent.at("id").is_string())
                {
                    agentIdStr = agent.at("id").template get<std::string>();
                }
                else
                {
                    agentIdStr = std::to_string(agent.at("id").template get<int>());
                }

                wdbSocketWrapper.query(WazuhDBQueryBuilder::builder().agentGetPackagesCommand(agentIdStr).build(),
                                       response);

                // Validate the response
                if (response.empty())
                {
                    logWarn(WM_VULNSCAN_LOGTAG,
                            "Empty response for agent '%s' in Wazuh-DB 'sys_programs' query",
                            agentIdStr.c_str());
                    continue;
                }

                for (const auto& package : response)
                {
                    flatbuffers::FlatBufferBuilder fbBuilder;
                    auto agentInfo =
                        SyscollectorSynchronization::CreateAgentInfoDirect(fbBuilder,
                                                                           agentIdStr.c_str(),
                                                                           agent.value("ip", "").c_str(),
                                                                           agent.value("name", "").c_str(),
                                                                           agent.value("version", "").c_str(),
                                                                           agent.value("node_name", "").c_str());
                    auto syscollectorPackage = SyscollectorSynchronization::Createsyscollector_packagesDirect(
                        fbBuilder,
                        package.value("architecture", "").c_str(),
                        package.value("checksum", "").c_str(),
                        package.value("description", "").c_str(),
                        package.value("format", "").c_str(),
                        package.value("groups", "").c_str(),
                        package.value("install_time", "").c_str(),
                        package.value("item_id", "").c_str(),
                        package.value("location", "").c_str(),
                        package.value("multiarch", "").c_str(),
                        package.value("name", "").c_str(),
                        package.value("priority", "").c_str(),
                        package.value("scan_time", "").c_str(),
                        package.value("size", 0),
                        package.value("source", "").c_str(),
                        package.value("vendor", "").c_str(),
                        package.value("version", "").c_str());

                    auto stateMsg = SyscollectorSynchronization::Createstate(
                        fbBuilder,
                        SyscollectorSynchronization::AttributesUnion::AttributesUnion_syscollector_packages,
                        syscollectorPackage.Union());

                    auto syncMsg = SyscollectorSynchronization::CreateSyncMsg(
                        fbBuilder,
                        agentInfo,
                        SyscollectorSynchronization::DataUnion::DataUnion_state,
                        stateMsg.Union());

                    fbBuilder.Finish(syncMsg);

                    std::variant<const SyscollectorDeltas::Delta*,
                                 const SyscollectorSynchronization::SyncMsg*,
                                 const nlohmann::json*>
                        variantData = SyscollectorSynchronization::GetSyncMsg(fbBuilder.GetBufferPointer());

                    // This could throw a exception if there is not information about the os, because the user
                    // disable it from syscollector.
                    try
                    {
                        auto context = std::make_shared<TScanContext>(variantData);
                        m_suborchestration->handleRequest(std::move(context));
                    }
                    catch (const std::exception& e)
                    {
                        logDebug1(WM_VULNSCAN_LOGTAG, "Unable to handle request: %s", e.what());
                    }
                }
            }
        }
        else
        {
            throw std::runtime_error("Invalid operation type in TQueryAllPackages::handleRequest");
        }
        return AbstractHandler<std::shared_ptr<TScanContext>>::handleRequest(std::move(data));
    }
};

using QueryAllPackages = TQueryAllPackages<>;

#endif // __QUERY_ALL_PKGS_FOR_AGT_HPP
